import pygame
import math
import random
import sys

# Initialize Pygame
pygame.init()

# Constants
SCREEN_WIDTH = 800
SCREEN_HEIGHT = 600
FPS = 60

# Colors
WHITE = (255, 255, 255)
BLACK = (0, 0, 0)
RED = (255, 0, 0)
GREEN = (0, 255, 0)
BLUE = (0, 0, 255)
YELLOW = (255, 255, 0)
GRAY = (128, 128, 128)
DARK_GRAY = (64, 64, 64)
ROAD_COLOR = (80, 80, 80)
GRASS_COLOR = (0, 150, 0)


class Car:
    def __init__(self, x, y, color, is_player=False):
        self.x = x
        self.y = y
        self.color = color
        self.speed = 0
        self.angle = 0
        self.max_speed = 8
        self.acceleration = 0.2
        self.friction = 0.1
        self.turn_speed = 3
        self.is_player = is_player
        self.width = 30
        self.height = 60

        # AI properties for non-player cars
        if not is_player:
            self.target_x = x
            self.ai_speed = random.uniform(2, 5)
            self.lane_change_timer = 0

    def update(self, keys_pressed=None):
        if self.is_player and keys_pressed:
            # Player controls
            if keys_pressed[pygame.K_UP] or keys_pressed[pygame.K_w]:
                self.speed = min(self.speed + self.acceleration, self.max_speed)
            elif keys_pressed[pygame.K_DOWN] or keys_pressed[pygame.K_s]:
                self.speed = max(self.speed - self.acceleration, -self.max_speed / 2)
            else:
                # Apply friction
                if self.speed > 0:
                    self.speed = max(0, self.speed - self.friction)
                else:
                    self.speed = min(0, self.speed + self.friction)

            if keys_pressed[pygame.K_LEFT] or keys_pressed[pygame.K_a]:
                if abs(self.speed) > 0.5:
                    self.angle -= self.turn_speed
            if keys_pressed[pygame.K_RIGHT] or keys_pressed[pygame.K_d]:
                if abs(self.speed) > 0.5:
                    self.angle += self.turn_speed
        else:
            # AI behavior
            self.speed = self.ai_speed

            # Simple lane changing behavior
            self.lane_change_timer += 1
            if self.lane_change_timer > random.randint(120, 300):
                self.target_x = random.randint(250, 550)
                self.lane_change_timer = 0

            # Move towards target lane
            if abs(self.x - self.target_x) > 5:
                if self.x < self.target_x:
                    self.x += 1
                else:
                    self.x -= 1

        # Update position based on angle and speed
        self.x += math.sin(math.radians(self.angle)) * self.speed
        self.y += math.cos(math.radians(self.angle)) * self.speed

        # Keep cars on screen (for player)
        if self.is_player:
            self.x = max(200, min(600, self.x))

    def draw(self, screen, camera_y):
        # Calculate screen position relative to camera
        screen_y = self.y - camera_y + SCREEN_HEIGHT // 2

        # Only draw if visible
        if -100 < screen_y < SCREEN_HEIGHT + 100:
            # Create car rectangle
            car_rect = pygame.Rect(0, 0, self.width, self.height)
            car_surface = pygame.Surface((self.width, self.height))
            car_surface.fill(self.color)

            # Add some detail to the car
            pygame.draw.rect(car_surface, WHITE, (5, 5, 20, 15))  # Windshield
            pygame.draw.rect(car_surface, WHITE, (5, 40, 20, 15))  # Rear window
            pygame.draw.circle(car_surface, BLACK, (8, 15), 4)  # Wheels
            pygame.draw.circle(car_surface, BLACK, (22, 15), 4)
            pygame.draw.circle(car_surface, BLACK, (8, 45), 4)
            pygame.draw.circle(car_surface, BLACK, (22, 45), 4)

            # Rotate the car
            rotated_car = pygame.transform.rotate(car_surface, -self.angle)
            rotated_rect = rotated_car.get_rect(center=(self.x, screen_y))

            screen.blit(rotated_car, rotated_rect)


class Road:
    def __init__(self):
        self.lane_width = 100
        self.road_width = 400
        self.road_x = SCREEN_WIDTH // 2 - self.road_width // 2

    def draw(self, screen, camera_y):
        # Draw grass background
        screen.fill(GRASS_COLOR)

        # Draw road
        road_rect = pygame.Rect(self.road_x, 0, self.road_width, SCREEN_HEIGHT)
        pygame.draw.rect(screen, ROAD_COLOR, road_rect)

        # Draw lane markings
        line_spacing = 40
        line_height = 20
        offset = int(camera_y % (line_spacing + line_height))

        # Center line
        center_x = self.road_x + self.road_width // 2
        for y in range(-offset, SCREEN_HEIGHT + line_height, line_spacing + line_height):
            pygame.draw.rect(screen, YELLOW, (center_x - 2, y, 4, line_height))

        # Side lines
        for y in range(-offset, SCREEN_HEIGHT + line_height, line_spacing + line_height):
            # Left lane line
            pygame.draw.rect(screen, WHITE, (self.road_x + self.lane_width - 2, y, 4, line_height))
            # Right lane line
            pygame.draw.rect(screen, WHITE, (self.road_x + self.road_width - self.lane_width - 2, y, 4, line_height))

        # Road edges
        pygame.draw.line(screen, WHITE, (self.road_x, 0), (self.road_x, SCREEN_HEIGHT), 4)
        pygame.draw.line(screen, WHITE, (self.road_x + self.road_width, 0),
                         (self.road_x + self.road_width, SCREEN_HEIGHT), 4)


class Game:
    def __init__(self):
        self.screen = pygame.display.set_mode((SCREEN_WIDTH, SCREEN_HEIGHT))
        pygame.display.set_caption("3D Car Racing Game")
        self.clock = pygame.time.Clock()

        # Create player car
        self.player = Car(SCREEN_WIDTH // 2, SCREEN_HEIGHT // 2, RED, True)

        # Create AI cars
        self.ai_cars = []
        for i in range(8):
            x = random.randint(250, 550)
            y = random.randint(-500, 500)
            color = random.choice([BLUE, GREEN, YELLOW, WHITE, BLACK])
            self.ai_cars.append(Car(x, y, color, False))

        self.road = Road()
        self.camera_y = 0
        self.score = 0
        self.font = pygame.font.Font(None, 36)
        self.small_font = pygame.font.Font(None, 24)

    def handle_events(self):
        for event in pygame.event.get():
            if event.type == pygame.QUIT:
                return False
            elif event.type == pygame.KEYDOWN:
                if event.key == pygame.K_ESCAPE:
                    return False
        return True

    def update(self):
        keys_pressed = pygame.key.get_pressed()

        # Update player
        self.player.update(keys_pressed)

        # Update camera to follow player
        self.camera_y = self.player.y - SCREEN_HEIGHT // 2

        # Update AI cars
        for car in self.ai_cars:
            car.update()

            # Reset AI cars that go too far behind
            if car.y < self.player.y - 600:
                car.y = self.player.y + random.randint(300, 800)
                car.x = random.randint(250, 550)

        # Update score based on distance traveled
        self.score = max(0, int(self.player.y / 10))

        # Spawn new AI cars ahead
        if random.randint(1, 100) == 1:
            if len(self.ai_cars) < 12:
                x = random.randint(250, 550)
                y = self.player.y + random.randint(400, 800)
                color = random.choice([BLUE, GREEN, YELLOW, WHITE, BLACK])
                self.ai_cars.append(Car(x, y, color, False))

        # Remove cars that are too far behind
        self.ai_cars = [car for car in self.ai_cars if car.y > self.player.y - 800]

    def draw(self):
        # Draw road
        self.road.draw(self.screen, self.camera_y)

        # Draw all cars
        # Sort cars by y position for proper depth
        all_cars = [self.player] + self.ai_cars
        all_cars.sort(key=lambda car: car.y)

        for car in all_cars:
            car.draw(self.screen, self.camera_y)

        # Draw UI
        score_text = self.font.render(f"Distance: {self.score}m", True, WHITE)
        speed_text = self.small_font.render(f"Speed: {abs(self.player.speed):.1f}", True, WHITE)
        controls_text = self.small_font.render("WASD/Arrow Keys to drive, ESC to quit", True, WHITE)

        self.screen.blit(score_text, (10, 10))
        self.screen.blit(speed_text, (10, 50))
        self.screen.blit(controls_text, (10, SCREEN_HEIGHT - 30))

        pygame.display.flip()

    def run(self):
        running = True
        while running:
            running = self.handle_events()
            self.update()
            self.draw()
            self.clock.tick(FPS)

        pygame.quit()
        sys.exit()


if __name__ == "__main__":
    game = Game()
    game.run()